from osgeo import gdal,ogr,osr,gdalconst
from numba import jit,njit
import numpy as np


@njit # 使用Numba的JIT编译器来加速函数
def connected8_2_connected4(in_arr):
    pattern_1 = np.array([[255, 0], [0, 255]], dtype=np.uint8)
    pattern_2 = np.array([[0, 255], [255, 0]], dtype=np.uint8)
    replacement = np.array([[255, 255], [255, 255]], dtype=np.uint8)
    out_arr = in_arr.copy()
    rows, cols = in_arr.shape # 行，列
    for i in range(rows-1):
        for j in range(cols-1):
            tmp_arr = out_arr[i:i+2,j:j+2].copy() 
            if np.array_equal(tmp_arr, pattern_1):
                out_arr[i:i+2,j:j+2] = replacement
            elif np.array_equal(tmp_arr, pattern_2):
                out_arr[i:i+2,j:j+2] = replacement
    if not np.array_equal(in_arr,out_arr):
        out_arr = connected8_2_connected4(out_arr)
    return out_arr


@njit # @jit(nopython=True)  # 使用Numba的JIT编译器来加速函数
def del_single_pixel(in_arr):
    pattern = np.array([[0, 0, 0], [0, 255, 0], [0, 0, 0]], dtype=np.uint8)
    replacement = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]], dtype=np.uint8)
    out_arr = in_arr.copy()
    rows, cols = in_arr.shape # 行，列
    # for i,j in product(range(rows-2),range(cols-2)):
    for i in range(rows-2):
        for j in range(cols-2):
            tmp_arr = out_arr[i:i+3,j:j+3].copy() 
            if np.array_equal(tmp_arr, pattern):
                out_arr[i:i+3,j:j+3] = replacement
    return out_arr


def raster_2_geometry(tif_path,vector_path='',drivername="Memory"):
    dataset = gdal.Open(tif_path)
    projection = dataset.GetProjection()
    srs = osr.SpatialReference(wkt=projection)
    # if srs.IsGeographic():print("使用的是地理坐标系")
    epsg_code = srs.GetAttrValue('AUTHORITY', 1)
    if epsg_code != '4326':
        mem_path = '/vsimem/output.tif'
        print("转WGS84坐标系")
        gdal.Warp(
            mem_path, 
            dataset, 
            dstSRS='EPSG:4326',
            resampleAlg=gdal.GRA_Bilinear,
            callback=gdal.TermProgress_nocb)
    else:
        print("该栅格使用的是WGS84坐标系。")
        mem_path = tif_path
    del dataset
    dataset = gdal.Open(mem_path) 
    xsize = dataset.RasterXSize
    ysize = dataset.RasterYSize
    # band_count = dataset.RasterCount
    geotransform = dataset.GetGeoTransform()
    projection = dataset.GetProjection()
    band = dataset.GetRasterBand(1)
    data_type = band.DataType
    # nodata = band.GetNoDataValue()
    mask_band = band.GetMaskBand()
    array = mask_band.ReadAsArray()
    del dataset 
    array = connected8_2_connected4(array)
    array = del_single_pixel(array)
    driver = gdal.GetDriverByName("MEM") # 'GTiff'
    mem_ds = driver.Create("", xsize, ysize, 1, data_type)
    mem_ds.SetGeoTransform(geotransform)  # 地理变换信息
    mem_ds.SetProjection(projection)  # 设置投影
    band = mem_ds.GetRasterBand(1)
    band.WriteArray(array)
    band.SetNoDataValue(0)
    mask_band = band.GetMaskBand()
    mem_ds.FlushCache()
    srs=osr.SpatialReference()
    srs.ImportFromEPSG(4326)
    driver = ogr.GetDriverByName(drivername) # 'ESRI Shapefile'
    vector_ds = driver.CreateDataSource(vector_path)
    layer = vector_ds.CreateLayer('layer',srs,geom_type=ogr.wkbPolygon)
    # layer.CreateField(ogr.FieldDefn("ID", ogr.OFTInteger))
    print("栅格转矢量")
    gdal.Polygonize(
        band, 
        mask_band, 
        layer, 
        -1, 
        options=[], 
        callback=gdal.TermProgress_nocb)
    del mem_ds
    featureCount = layer.GetFeatureCount()
    if featureCount == 1:
        feature = layer.GetFeature(0)
        reslut = geometry = feature.geometry()
    else:
        multi_polygon = ogr.Geometry(ogr.wkbMultiPolygon)
        feature_defn = layer.GetLayerDefn()
        for index in range(featureCount):
            feature = layer.GetFeature(index)
            geometry = feature.geometry()
            layer.DeleteFeature(index)
            multi_polygon.AddGeometry(geometry)
        feature = ogr.Feature(feature_defn)
        feature.SetGeometry(multi_polygon)
        layer.CreateFeature(feature)
        # layer.SyncToDisk()
        reslut = multi_polygon.ExportToWkt()
    vector_ds.Destroy()
    return reslut


# tif_path = r"/data/fengyy/dateset/temp/202403/131123_武强县20230726.tif"
# shp_path = r"/data/fengyy/dateset/temp/202403/131123_武强县20230726-01.shp"
tif_path = r"/data/fengyy/dateset/temp/202403/130533_威县20230507.tif"
shp_path = r"/data/fengyy/dateset/temp/202403/130533_威县20230507.shp"
raster_2_geometry(tif_path,shp_path,drivername='ESRI Shapefile')



